/*
 * Copyright 2025-2025 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.springframework.ai.mcp.server.autoconfigure;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.stream.Collectors;

import io.modelcontextprotocol.server.McpAsyncServer;
import io.modelcontextprotocol.server.McpAsyncServerExchange;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpServer.AsyncSpecification;
import io.modelcontextprotocol.server.McpServer.SyncSpecification;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.McpServerFeatures.AsyncCompletionSpecification;
import io.modelcontextprotocol.server.McpServerFeatures.AsyncPromptSpecification;
import io.modelcontextprotocol.server.McpServerFeatures.AsyncResourceSpecification;
import io.modelcontextprotocol.server.McpServerFeatures.AsyncToolSpecification;
import io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification;
import io.modelcontextprotocol.server.McpServerFeatures.SyncPromptSpecification;
import io.modelcontextprotocol.server.McpServerFeatures.SyncResourceSpecification;
import io.modelcontextprotocol.server.McpServerFeatures.SyncToolSpecification;
import io.modelcontextprotocol.server.McpSyncServer;
import io.modelcontextprotocol.server.McpSyncServerExchange;
import io.modelcontextprotocol.server.transport.StdioServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.Implementation;
import io.modelcontextprotocol.spec.McpServerTransportProvider;
import io.modelcontextprotocol.spec.McpServerTransportProviderBase;
import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider;
import reactor.core.publisher.Mono;

import org.springframework.ai.mcp.McpToolUtils;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.core.log.LogAccessor;
import org.springframework.util.CollectionUtils;
import org.springframework.util.MimeType;

/**
 * {@link EnableAutoConfiguration Auto-configuration} for the Model Context Protocol (MCP)
 * Server.
 * <p>
 * This configuration class sets up the core MCP server components with support for both
 * synchronous and asynchronous operation modes. The server type is controlled through the
 * {@code spring.ai.mcp.server.type} property, defaulting to SYNC mode.
 * <p>
 * Core features and capabilities include:
 * <ul>
 * <li>Tools: Extensible tool registration system supporting both sync and async
 * execution</li>
 * <li>Resources: Static and dynamic resource management with optional change
 * notifications</li>
 * <li>Prompts: Configurable prompt templates with change notification support</li>
 * <li>Transport: Flexible transport layer with built-in support for:
 * <ul>
 * <li>STDIO (default): Standard input/output based communication</li>
 * <li>WebMvc: HTTP-based transport when Spring MVC is available</li>
 * <li>WebFlux: Reactive transport when Spring WebFlux is available</li>
 * <li>Streamable: High-performance streaming transport</li>
 * </ul>
 * </li>
 * </ul>
 * <p>
 * The configuration is activated when:
 * <ul>
 * <li>The required MCP classes ({@link McpSchema} and {@link McpSyncServer}) are on the
 * classpath</li>
 * <li>The {@code spring.ai.mcp.server.enabled} property is true (default)</li>
 * </ul>
 * <p>
 * Server configuration is managed through {@link McpServerProperties} with support for:
 * <ul>
 * <li>Server identification (name, version)</li>
 * <li>Transport selection (STDIO, SSE, STREAMABLE)</li>
 * <li>Change notification settings for tools, resources, and prompts</li>
 * <li>Sync/Async operation mode selection</li>
 * </ul>
 * <p>
 * Web transport support is provided separately by specialized auto-configuration classes:
 * <ul>
 * <li>SSE transport: {@link McpWebMvcServerAutoConfiguration}, {@link McpWebFluxServerAutoConfiguration}</li>
 * <li>Streamable transport: {@link McpWebMvcStreamableServerAutoConfiguration},
 * {@link McpWebFluxStreamableServerAutoConfiguration}, {@link McpHttpServletStreamableServerAutoConfiguration}</li>
 * </ul>
 *
 * @author Christian Tzolov
 * @since 1.0.0
 * @see McpServerProperties
 * @see McpWebMvcServerAutoConfiguration
 * @see McpWebFluxServerAutoConfiguration
 * @see McpWebMvcStreamableServerAutoConfiguration
 * @see McpWebFluxStreamableServerAutoConfiguration
 * @see McpHttpServletStreamableServerAutoConfiguration
 * @see ToolCallback
 */
@AutoConfiguration(after = {
		McpWebMvcServerAutoConfiguration.class,
		McpWebFluxServerAutoConfiguration.class,
		McpWebMvcStreamableServerAutoConfiguration.class,
		McpWebFluxStreamableServerAutoConfiguration.class,
		McpHttpServletStreamableServerAutoConfiguration.class
})
@ConditionalOnClass({ McpSchema.class, McpSyncServer.class })
@EnableConfigurationProperties(McpServerProperties.class)
@ConditionalOnProperty(prefix = McpServerProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true",
		matchIfMissing = true)
public class McpServerAutoConfiguration {

	private static final LogAccessor logger = new LogAccessor(McpServerAutoConfiguration.class);

	/**
	 * STDIO传输提供者Bean
	 * <p>
	 * 当启用STDIO传输时创建，或者作为Web传输的fallback。
	 */
	@Bean
	@ConditionalOnMissingBean
	@ConditionalOnProperty(prefix = McpServerProperties.CONFIG_PREFIX, name = "stdio", havingValue = "true")
	public McpServerTransportProviderBase stdioServerTransport() {
		logger.info("创建STDIO传输提供者");
		return new StdioServerTransportProvider();
	}

	/**
	 * 服务器能力构建器
	 */
	@Bean
	@ConditionalOnMissingBean
	public McpSchema.ServerCapabilities.Builder capabilitiesBuilder() {
		return McpSchema.ServerCapabilities.builder();
	}

	/**
	 * 同步工具规范列表
	 * <p>
	 * 收集所有可用的工具回调并转换为同步工具规范。
	 */
	@Bean
	@ConditionalOnProperty(prefix = McpServerProperties.CONFIG_PREFIX, name = "type", havingValue = "SYNC",
			matchIfMissing = true)
	public List<McpServerFeatures.SyncToolSpecification> syncTools(ObjectProvider<List<ToolCallback>> toolCalls,
																   List<ToolCallback> toolCallbacksList, McpServerProperties serverProperties) {

		List<ToolCallback> tools = new ArrayList<>(toolCalls.stream().flatMap(List::stream).toList());

		if (!CollectionUtils.isEmpty(toolCallbacksList)) {
			tools.addAll(toolCallbacksList);
		}

		return this.toSyncToolSpecifications(tools, serverProperties);
	}

	private List<McpServerFeatures.SyncToolSpecification> toSyncToolSpecifications(List<ToolCallback> tools,
																				   McpServerProperties serverProperties) {

		// De-duplicate tools by their name, keeping the first occurrence of each tool name
		return tools.stream()
				.collect(Collectors.toMap(tool -> tool.getToolDefinition().name(), tool -> tool,
						(existing, _) -> existing))
				.values()
				.stream()
				.map(tool -> {
					String toolName = tool.getToolDefinition().name();
					MimeType mimeType = (serverProperties.getToolResponseMimeType().containsKey(toolName))
							? MimeType.valueOf(serverProperties.getToolResponseMimeType().get(toolName)) : null;
					return McpToolUtils.toSyncToolSpecification(tool, mimeType);
				})
				.toList();
	}

	/**
	 * 同步MCP服务器
	 * <p>
	 * 创建同步模式的MCP服务器，使用可用的传输提供者。
	 */
	@Bean
	@ConditionalOnProperty(prefix = McpServerProperties.CONFIG_PREFIX, name = "type", havingValue = "SYNC",
			matchIfMissing = true)
	@ConditionalOnMissingBean(McpSyncServer.class)
	public McpSyncServer mcpSyncServer(McpServerTransportProviderBase transportProvider,
									   McpSchema.ServerCapabilities.Builder capabilitiesBuilder, McpServerProperties serverProperties,
									   ObjectProvider<List<SyncToolSpecification>> tools,
									   ObjectProvider<List<SyncResourceSpecification>> resources,
									   ObjectProvider<List<SyncPromptSpecification>> prompts,
									   ObjectProvider<List<SyncCompletionSpecification>> completions,
									   ObjectProvider<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeConsumers,
									   List<ToolCallbackProvider> toolCallbackProvider) {

		logger.info("创建MCP同步服务器，传输: " + transportProvider.getClass().getSimpleName());

		McpSchema.Implementation serverInfo = new Implementation(serverProperties.getName(),
				serverProperties.getVersion());

		// Create the server with transport provider
		SyncSpecification<?> serverBuilder;
		if (transportProvider instanceof McpStreamableServerTransportProvider) {
			serverBuilder = McpServer.sync((McpStreamableServerTransportProvider) transportProvider);

		} else {
			serverBuilder = McpServer.sync((McpServerTransportProvider) transportProvider);
		}
		serverBuilder.serverInfo(serverInfo);
		// Tools
		if (serverProperties.getCapabilities().isTool()) {
			logger.info("启用工具能力，变更通知: " + serverProperties.isToolChangeNotification());
			capabilitiesBuilder.tools(serverProperties.isToolChangeNotification());

			List<SyncToolSpecification> toolSpecifications = new ArrayList<>(
					tools.stream().flatMap(List::stream).toList());

			List<ToolCallback> providerToolCallbacks = toolCallbackProvider.stream()
					.map(pr -> List.of(pr.getToolCallbacks()))
					.flatMap(List::stream)
					.filter(Objects::nonNull)
					.toList();

			toolSpecifications.addAll(this.toSyncToolSpecifications(providerToolCallbacks, serverProperties));

			if (!CollectionUtils.isEmpty(toolSpecifications)) {
				serverBuilder.tools(toolSpecifications);
				logger.info("注册工具数量: " + toolSpecifications.size());
			}
		}

		// Resources
		if (serverProperties.getCapabilities().isResource()) {
			logger.info("启用资源能力，变更通知: " + serverProperties.isResourceChangeNotification());
			capabilitiesBuilder.resources(false, serverProperties.isResourceChangeNotification());

			List<SyncResourceSpecification> resourceSpecifications = resources.stream().flatMap(List::stream).toList();
			if (!CollectionUtils.isEmpty(resourceSpecifications)) {
				serverBuilder.resources(resourceSpecifications);
				logger.info("注册资源数量: " + resourceSpecifications.size());
			}
		}

		// Prompts
		if (serverProperties.getCapabilities().isPrompt()) {
			logger.info("启用提示能力，变更通知: " + serverProperties.isPromptChangeNotification());
			capabilitiesBuilder.prompts(serverProperties.isPromptChangeNotification());

			List<SyncPromptSpecification> promptSpecifications = prompts.stream().flatMap(List::stream).toList();
			if (!CollectionUtils.isEmpty(promptSpecifications)) {
				serverBuilder.prompts(promptSpecifications);
				logger.info("注册提示数量: " + promptSpecifications.size());
			}
		}

		// Completions
		if (serverProperties.getCapabilities().isCompletion()) {
			logger.info("启用补全能力");
			capabilitiesBuilder.completions();

			List<SyncCompletionSpecification> completionSpecifications = completions.stream()
					.flatMap(List::stream)
					.toList();
			if (!CollectionUtils.isEmpty(completionSpecifications)) {
				serverBuilder.completions(completionSpecifications);
				logger.info("注册补全数量: " + completionSpecifications.size());
			}
		}

		rootsChangeConsumers.ifAvailable(consumer -> {
			serverBuilder.rootsChangeHandler((exchange, roots) -> consumer.accept(exchange, roots));
			logger.info("注册根变更消费者");
		});

		serverBuilder.capabilities(capabilitiesBuilder.build());
		serverBuilder.instructions(serverProperties.getInstructions());
		serverBuilder.requestTimeout(serverProperties.getRequestTimeout());

		McpSyncServer server = serverBuilder.build();
		logger.info("MCP同步服务器创建完成");
		return server;
	}

	/**
	 * 异步工具规范列表
	 * <p>
	 * 收集所有可用的工具回调并转换为异步工具规范。
	 */
	@Bean
	@ConditionalOnProperty(prefix = McpServerProperties.CONFIG_PREFIX, name = "type", havingValue = "ASYNC")
	public List<McpServerFeatures.AsyncToolSpecification> asyncTools(ObjectProvider<List<ToolCallback>> toolCalls,
																	 List<ToolCallback> toolCallbackList, McpServerProperties serverProperties) {

		List<ToolCallback> tools = new ArrayList<>(toolCalls.stream().flatMap(List::stream).toList());
		if (!CollectionUtils.isEmpty(toolCallbackList)) {
			tools.addAll(toolCallbackList);
		}

		return this.toAsyncToolSpecification(tools, serverProperties);
	}

	private List<McpServerFeatures.AsyncToolSpecification> toAsyncToolSpecification(List<ToolCallback> tools,
																					McpServerProperties serverProperties) {
		// De-duplicate tools by their name, keeping the first occurrence of each tool name
		return tools.stream()
				.collect(Collectors.toMap(tool -> tool.getToolDefinition().name(), tool -> tool,
						(existing, replacement) -> existing))
				.values()
				.stream()
				.map(tool -> {
					String toolName = tool.getToolDefinition().name();
					MimeType mimeType = (serverProperties.getToolResponseMimeType().containsKey(toolName))
							? MimeType.valueOf(serverProperties.getToolResponseMimeType().get(toolName)) : null;
					return McpToolUtils.toAsyncToolSpecification(tool, mimeType);
				})
				.toList();
	}

	/**
	 * 异步MCP服务器
	 * <p>
	 * 创建异步模式的MCP服务器，使用可用的传输提供者。
	 */
	@Bean
	@ConditionalOnProperty(prefix = McpServerProperties.CONFIG_PREFIX, name = "type", havingValue = "ASYNC")
	@ConditionalOnMissingBean(McpAsyncServer.class)
	public McpAsyncServer mcpAsyncServer(McpServerTransportProviderBase transportProvider,
										 McpSchema.ServerCapabilities.Builder capabilitiesBuilder, McpServerProperties serverProperties,
										 ObjectProvider<List<AsyncToolSpecification>> tools,
										 ObjectProvider<List<AsyncResourceSpecification>> resources,
										 ObjectProvider<List<AsyncPromptSpecification>> prompts,
										 ObjectProvider<List<AsyncCompletionSpecification>> completions,
										 ObjectProvider<BiConsumer<McpAsyncServerExchange, List<McpSchema.Root>>> rootsChangeConsumer,
										 List<ToolCallbackProvider> toolCallbackProvider) {

		logger.info("创建MCP异步服务器，传输: " + transportProvider.getClass().getSimpleName());

		McpSchema.Implementation serverInfo = new Implementation(serverProperties.getName(),
				serverProperties.getVersion());

		// Create the server with transport provider
		AsyncSpecification<?> serverBuilder;
		if (transportProvider instanceof McpStreamableServerTransportProvider) {
			serverBuilder = McpServer.async((McpStreamableServerTransportProvider) transportProvider);

		} else {
			serverBuilder = McpServer.async((McpServerTransportProvider) transportProvider);
		}
		serverBuilder.serverInfo(serverInfo);
		// Tools
		if (serverProperties.getCapabilities().isTool()) {
			List<AsyncToolSpecification> toolSpecifications = new ArrayList<>(
					tools.stream().flatMap(List::stream).toList());
			List<ToolCallback> providerToolCallbacks = toolCallbackProvider.stream()
					.map(pr -> List.of(pr.getToolCallbacks()))
					.flatMap(List::stream)
					.filter(fc -> fc instanceof ToolCallback)
					.map(fc -> (ToolCallback) fc)
					.toList();

			toolSpecifications.addAll(this.toAsyncToolSpecification(providerToolCallbacks, serverProperties));

			logger.info("启用异步工具能力，变更通知: " + serverProperties.isToolChangeNotification());
			capabilitiesBuilder.tools(serverProperties.isToolChangeNotification());

			if (!CollectionUtils.isEmpty(toolSpecifications)) {
				serverBuilder.tools(toolSpecifications);
				logger.info("注册异步工具数量: " + toolSpecifications.size());
			}
		}

		// Resources
		if (serverProperties.getCapabilities().isResource()) {
			logger.info("启用异步资源能力，变更通知: " + serverProperties.isResourceChangeNotification());
			capabilitiesBuilder.resources(false, serverProperties.isResourceChangeNotification());

			List<AsyncResourceSpecification> resourceSpecifications = resources.stream().flatMap(List::stream).toList();
			if (!CollectionUtils.isEmpty(resourceSpecifications)) {
				serverBuilder.resources(resourceSpecifications);
				logger.info("注册异步资源数量: " + resourceSpecifications.size());
			}
		}

		// Prompts
		if (serverProperties.getCapabilities().isPrompt()) {
			logger.info("启用异步提示能力，变更通知: " + serverProperties.isPromptChangeNotification());
			capabilitiesBuilder.prompts(serverProperties.isPromptChangeNotification());
			List<AsyncPromptSpecification> promptSpecifications = prompts.stream().flatMap(List::stream).toList();

			if (!CollectionUtils.isEmpty(promptSpecifications)) {
				serverBuilder.prompts(promptSpecifications);
				logger.info("注册异步提示数量: " + promptSpecifications.size());
			}
		}

		// Completions
		if (serverProperties.getCapabilities().isCompletion()) {
			logger.info("启用异步补全能力");
			capabilitiesBuilder.completions();
			List<AsyncCompletionSpecification> completionSpecifications = completions.stream()
					.flatMap(List::stream)
					.toList();

			if (!CollectionUtils.isEmpty(completionSpecifications)) {
				serverBuilder.completions(completionSpecifications);
				logger.info("注册异步补全数量: " + completionSpecifications.size());
			}
		}

		rootsChangeConsumer.ifAvailable(consumer -> {
			BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>> asyncConsumer = (exchange, roots) -> {
				consumer.accept(exchange, roots);
				return Mono.empty();
			};
			serverBuilder.rootsChangeHandler(asyncConsumer);
			logger.info("注册异步根变更消费者");
		});

		serverBuilder.capabilities(capabilitiesBuilder.build());
		serverBuilder.instructions(serverProperties.getInstructions());
		serverBuilder.requestTimeout(serverProperties.getRequestTimeout());

		McpAsyncServer server = serverBuilder.build();
		logger.info("MCP异步服务器创建完成");
		return server;
	}

}